from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
import imageio
from PIL import Image
import numpy as np
import os
import time
from loguru import logger

body_part = {
    "Face": [],
    "Hair": [],
    "Torso-skin": [],
    "UpperClothes": [],
    "Left-arm": [],
    "Right-arm": [],
    "Coat": [],
    "Sunglasses": [],
    "Scarf": [],
    "Skirt": [],
    "Pants": [],
    "Dress": [],
    "Right-shoe": [],
    "Right-leg": [],
    "Left-shoe": [],
    "Left-leg": [],
}


class BodySegModel:
    def __init__(self):
        time_before = time.time()
        self.model = pipeline(Tasks.image_segmentation,
                              "damo/cv_resnet101_image-multiple-human-parsing")
        time_after = time.time()
        logger.info(f"[BodySeg]模型加载时间: {time_after - time_before}")


class BodySeg:
    color_white = [255, 255, 255]
    color_black = [0, 0, 0]

    def __init__(self, seg_target=["Face", "Hair"], model: BodySegModel = None):
        self.model = model.model if model else BodySegModel().model
        self.canvas = None
        self.body = dict(body_part)
        # self.single_body = ["Face","Torso-skin"]
        self.single_body = []
        self.seg_target = seg_target
        self.seg_color = self.color_black
        self.remove_color = self.color_white

    @classmethod
    def _is_target_(cls, target, x, y):
        r, g, b = target.getpixel((x, y))
        # 如果是白色，则说明是目标区域
        return r == 255 and g == 255 and b == 255

    @classmethod
    def _change_bg_color_(cls, img, x, y, colors):
        img.putpixel((x, y), (colors[0], colors[1], colors[2]))

    @classmethod
    def _get_item_(cls, masks, labels, label):
        if label in labels:
            iis = [i for i, x in enumerate(labels) if x == label]
            imgs = [None] * len(iis)
            for i, v in enumerate(iis):
                imgs[i] = Image.fromarray(masks[v] * 255).convert("RGB")
            return imgs
        else:
            return []

    def _get_body_(self, infer_result):
        labels = infer_result[OutputKeys.LABELS]
        logger.info(f'获取到人体部位：{labels}')
        masks = infer_result['masks']
        for k, v in self.body.items():
            if k in self.single_body:
                self.body[k] = Image.fromarray(
                    masks[labels.index(k)] * 255).convert("RGB")
            else:
                self.body[k] = self._get_item_(masks, labels, k)

    def _set_mask_(self, targets, x, y, color, multi=True):
        if multi:
            for target in targets:
                if self._is_target_(target, x, y):
                    self._change_bg_color_(self.canvas, x, y, color)
        else:
            if self._is_target_(targets, x, y):
                self._change_bg_color_(self.canvas, x, y, color)

    def get_mask(self, image_path, save_dir):
        time_before = time.time()
        result = self.model(image_path)
        time_after = time.time()
        logger.info(f"模型推理时间: {time_after - time_before}")
        self._get_body_(infer_result=result)
        face = self.body["Face"][0]
        # 创建画布底图,白色底
        self.canvas = Image.new(
            "RGB", (face.width, face.height), tuple(self.remove_color))
        for x in range(face.width):
            for y in range(face.height):
                for k, v in self.body.items():
                    if k in self.single_body:
                        self._set_mask_(targets=self.body[k], x=x, y=y,
                                        color=self.seg_color if k in self.seg_target else self.remove_color,
                                        multi=False)
                    else:
                        self._set_mask_(targets=self.body[k], x=x, y=y,
                                        color=self.seg_color if k in self.seg_target else self.remove_color,
                                        multi=True)

        filename = os.path.join(
            save_dir, "body-mask-" + str(time.time()) + ".png")
        self.canvas.save(filename)
        return filename


class NeckBodySeg(BodySeg):
    def __init__(self, model: BodySegModel = None):
        super().__init__(seg_target=["Torso-skin"], model=model)


class HairBodySeg(BodySeg):
    def __init__(self, model: BodySegModel = None):
        super().__init__(seg_target=["Hair"], model=model)


class ClothBodySeg(BodySeg):
    def __init__(self, model: BodySegModel = None):
        super().__init__(seg_target=[
            "Coat", "Dress", "UpperClothes"], model=model)
